import math
import torch
import torch.nn as nn
from torchvision import models


class ResNetBYOL(nn.Module):
    """Teacher model.

    Bootstrap Your Own Latent(BYOL):
    Reference: https://medium.com/swlh/neural-networks-intuitions-10-byol-paper-explanation-f8b1d6e83b1c

    """
    def __init__(self, projection_in_dimension: int, projection_out_dimension: int, num_class: int, mode: str) -> None:
        super().__init__()
        backbone = models.resnet50(pretrained=True)
        # backbone = models.resnext50_32x4d(pretrained=True)
        # backbone = pretrainedmodels.__dict__['inceptionresnetv2'](num_classes=1000, pretrained='imagenet')
        self.encoder = nn.Sequential(*list(backbone.children())[:-1])
        self.encoder[0] = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.projection = nn.Sequential(nn.Linear(backbone.fc.in_features, projection_in_dimension),
                                        nn.BatchNorm1d(projection_in_dimension), nn.ReLU(),
                                        nn.Linear(projection_in_dimension, projection_out_dimension))
        if mode == 'online':
            self.cls = nn.Sequential(nn.Linear(backbone.fc.in_features, num_class))

    def forward(self, x: torch.tensor, return_cls: bool = False, return_features: bool = False):
        features = self.encoder(x)
        features = features.view(features.shape[0], features.shape[1])
        if return_cls:
            return self.projection(features), self.cls(features)
        elif return_features:
            return self.projection(features), features
        return self.projection(features)


class MLPHead(nn.Module):
    """BYOL MLP head."""
    def __init__(self, input_channel, projection_in_dimension, projection_out_dimension) -> None:
        super().__init__()
        self.mlp_head = nn.Sequential(nn.Linear(input_channel, projection_in_dimension),
                                      nn.BatchNorm1d(projection_in_dimension), nn.ReLU(),
                                      nn.Linear(projection_in_dimension, projection_out_dimension))

    def forward(self, x):
        return self.mlp_head(x)


class ResNextCls(nn.Module):
    """Student model."""
    def __init__(self, num_class: int) -> None:
        super().__init__()
        backbone = models.resnet50(pretrained=True)
        # backbone = models.resnext50_32x4d(pretrained=True)
        # backbone = pretrainedmodels.__dict__['inceptionresnetv2'](num_classes=1000, pretrained='imagenet')
        self.encoder = nn.Sequential(*list(backbone.children())[:-1])
        self.encoder[0] = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.cls = nn.Sequential(nn.Linear(2048, num_class))

    def forward(self, x: torch.tensor, return_features: bool = False):
        features = self.encoder(x)
        features = features.view(features.shape[0], features.shape[1])
        if return_features:
            return self.cls(features), features
        return self.cls(features)


class MetricNet(nn.Module):
    """Metric model."""
    def __init__(self) -> None:
        super().__init__()
        self.model = nn.Sequential(nn.Linear(2048, 1024), nn.Linear(1024, 256))

    def forward(self, x: torch.tensor):
        features = self.model(x)
        return features
